"""
进度跟踪器
实现PPT生成过程的进度跟踪和状态管理，提供实时状态更新。
"""

import logging
import asyncio
from typing import Dict, List, Optional, Any, Callable
from datetime import datetime
from enum import Enum

from ..models import GenerationSession, GenerationStatus, GenerationStep
from ..services.websocket_manager import connection_manager

logger = logging.getLogger(__name__)


class ProgressStage(str, Enum):
    """进度阶段枚举"""
    INITIALIZING = "initializing"
    GENERATING_OUTLINE = "generating_outline"
    CREATING_CONTENT = "creating_content"
    APPLYING_TEMPLATE = "applying_template"
    GENERATING_PPT = "generating_ppt"
    FINALIZING = "finalizing"
    COMPLETED = "completed"
    ERROR = "error"


class ProgressTracker:
    """
    进度跟踪器类
    
    负责管理PPT生成过程的进度跟踪和状态更新
    """
    
    def __init__(self):
        """初始化进度跟踪器"""
        # 存储活跃的生成会话
        self.active_sessions: Dict[str, GenerationSession] = {}
        
        # 进度阶段配置
        self.stage_weights = {
            ProgressStage.INITIALIZING: 5,
            ProgressStage.GENERATING_OUTLINE: 20,
            ProgressStage.CREATING_CONTENT: 40,
            ProgressStage.APPLYING_TEMPLATE: 15,
            ProgressStage.GENERATING_PPT: 15,
            ProgressStage.FINALIZING: 5
        }
        
        # 进度阶段描述
        self.stage_descriptions = {
            ProgressStage.INITIALIZING: "初始化生成环境...",
            ProgressStage.GENERATING_OUTLINE: "AI正在生成演示文稿大纲...",
            ProgressStage.CREATING_CONTENT: "AI正在创建幻灯片内容...",
            ProgressStage.APPLYING_TEMPLATE: "应用模板样式...",
            ProgressStage.GENERATING_PPT: "生成PPT文件...",
            ProgressStage.FINALIZING: "完成最终处理...",
            ProgressStage.COMPLETED: "生成完成！",
            ProgressStage.ERROR: "生成过程中出现错误"
        }
    
    async def start_session(
        self, 
        session_id: str, 
        total_slides: int = 10,
        template_id: str = "professional"
    ) -> GenerationSession:
        """
        开始新的生成会话
        
        Args:
            session_id: 会话ID
            total_slides: 总幻灯片数量
            template_id: 模板ID
            
        Returns:
            GenerationSession: 生成会话对象
        """
        try:
            # 创建新的生成会话
            session = GenerationSession(
                session_id=session_id,
                status=GenerationStatus.PENDING,
                progress=0,
                current_step=GenerationStep.INITIALIZING,
                template_config={"template_id": template_id},
                created_at=datetime.now(),
                updated_at=datetime.now()
            )
            
            # 存储会话
            self.active_sessions[session_id] = session
            
            # 发送初始状态
            await self._send_progress_update(session_id, session)
            
            logger.info(f"开始生成会话: {session_id}")
            return session
            
        except Exception as e:
            logger.error(f"开始生成会话失败: {e}")
            raise
    
    async def update_progress(
        self,
        session_id: str,
        stage: ProgressStage,
        stage_progress: float = 0.0,
        message: Optional[str] = None,
        details: Optional[Dict[str, Any]] = None
    ):
        """
        更新生成进度
        
        Args:
            session_id: 会话ID
            stage: 当前阶段
            stage_progress: 阶段内进度 (0.0-1.0)
            message: 自定义消息
            details: 额外详情
        """
        if session_id not in self.active_sessions:
            logger.warning(f"会话不存在: {session_id}")
            return
        
        try:
            session = self.active_sessions[session_id]
            
            # 计算总体进度
            total_progress = self._calculate_total_progress(stage, stage_progress)
            
            # 更新会话状态
            session.progress = min(100, max(0, int(total_progress)))
            session.current_step = self._stage_to_step(stage)
            session.status = GenerationStatus.GENERATING
            session.updated_at = datetime.now()
            
            # 准备进度消息
            progress_message = message or self.stage_descriptions.get(stage, f"处理中: {stage}")
            
            # 发送进度更新
            await self._send_progress_update(
                session_id, 
                session, 
                progress_message,
                details
            )
            
            logger.debug(f"进度更新 - 会话: {session_id}, 阶段: {stage}, 进度: {total_progress:.1f}%")
            
        except Exception as e:
            logger.error(f"更新进度失败: {e}")
            await self.set_error(session_id, f"进度更新失败: {str(e)}")
    
    async def update_slide_progress(
        self,
        session_id: str,
        current_slide: int,
        total_slides: int,
        slide_title: str = ""
    ):
        """
        更新幻灯片生成进度
        
        Args:
            session_id: 会话ID
            current_slide: 当前幻灯片编号
            total_slides: 总幻灯片数量
            slide_title: 幻灯片标题
        """
        if total_slides <= 0:
            return
        
        # 计算幻灯片进度
        slide_progress = current_slide / total_slides
        
        # 构建消息
        message = f"正在生成第 {current_slide}/{total_slides} 张幻灯片"
        if slide_title:
            message += f": {slide_title}"
        
        # 更新进度
        await self.update_progress(
            session_id=session_id,
            stage=ProgressStage.CREATING_CONTENT,
            stage_progress=slide_progress,
            message=message,
            details={
                "current_slide": current_slide,
                "total_slides": total_slides,
                "slide_title": slide_title
            }
        )
    
    async def complete_session(
        self,
        session_id: str,
        file_path: str,
        slides_count: int
    ):
        """
        完成生成会话
        
        Args:
            session_id: 会话ID
            file_path: 生成的文件路径
            slides_count: 生成的幻灯片数量
        """
        if session_id not in self.active_sessions:
            logger.warning(f"会话不存在: {session_id}")
            return
        
        try:
            session = self.active_sessions[session_id]
            
            # 更新会话状态
            session.status = GenerationStatus.COMPLETED
            session.progress = 100
            session.current_step = GenerationStep.COMPLETED
            session.file_path = file_path
            session.updated_at = datetime.now()
            
            # 发送完成消息
            await connection_manager.send_completion_message(
                session_id=session_id,
                file_path=file_path,
                slides_count=slides_count
            )
            
            # 发送最终进度更新
            await self._send_progress_update(
                session_id,
                session,
                "PPT生成完成！",
                {
                    "file_path": file_path,
                    "slides_count": slides_count,
                    "completion_time": datetime.now().isoformat()
                }
            )
            
            logger.info(f"生成会话完成: {session_id}, 文件: {file_path}")
            
        except Exception as e:
            logger.error(f"完成生成会话失败: {e}")
            await self.set_error(session_id, f"完成会话失败: {str(e)}")
    
    async def set_error(
        self,
        session_id: str,
        error_message: str,
        error_code: Optional[str] = None
    ):
        """
        设置会话错误状态
        
        Args:
            session_id: 会话ID
            error_message: 错误消息
            error_code: 错误代码
        """
        if session_id not in self.active_sessions:
            logger.warning(f"会话不存在: {session_id}")
            return
        
        try:
            session = self.active_sessions[session_id]
            
            # 更新会话状态
            session.status = GenerationStatus.ERROR
            session.error_message = error_message
            session.updated_at = datetime.now()
            
            # 发送错误消息
            await connection_manager.send_error_message(
                session_id=session_id,
                error=error_message,
                error_code=error_code
            )
            
            # 发送错误进度更新
            await self._send_progress_update(
                session_id,
                session,
                f"生成失败: {error_message}",
                {
                    "error_code": error_code,
                    "error_time": datetime.now().isoformat()
                }
            )
            
            logger.error(f"生成会话错误: {session_id}, 错误: {error_message}")
            
        except Exception as e:
            logger.error(f"设置错误状态失败: {e}")
    
    async def cancel_session(self, session_id: str):
        """
        取消生成会话
        
        Args:
            session_id: 会话ID
        """
        if session_id not in self.active_sessions:
            logger.warning(f"会话不存在: {session_id}")
            return
        
        try:
            session = self.active_sessions[session_id]
            
            # 更新会话状态
            session.status = GenerationStatus.CANCELLED
            session.updated_at = datetime.now()
            
            # 发送取消消息
            await self._send_progress_update(
                session_id,
                session,
                "生成已取消",
                {"cancelled_at": datetime.now().isoformat()}
            )
            
            logger.info(f"生成会话已取消: {session_id}")
            
        except Exception as e:
            logger.error(f"取消生成会话失败: {e}")
    
    def get_session(self, session_id: str) -> Optional[GenerationSession]:
        """
        获取生成会话
        
        Args:
            session_id: 会话ID
            
        Returns:
            GenerationSession: 会话对象，如果不存在则返回None
        """
        return self.active_sessions.get(session_id)
    
    def get_all_sessions(self) -> Dict[str, GenerationSession]:
        """
        获取所有活跃会话
        
        Returns:
            Dict: 所有活跃会话的字典
        """
        return self.active_sessions.copy()
    
    def cleanup_session(self, session_id: str):
        """
        清理会话数据
        
        Args:
            session_id: 会话ID
        """
        if session_id in self.active_sessions:
            del self.active_sessions[session_id]
            logger.info(f"清理会话数据: {session_id}")
    
    def cleanup_old_sessions(self, max_age_hours: int = 24):
        """
        清理旧的会话数据
        
        Args:
            max_age_hours: 最大保留时间（小时）
        """
        try:
            current_time = datetime.now()
            old_sessions = []
            
            for session_id, session in self.active_sessions.items():
                age_hours = (current_time - session.created_at).total_seconds() / 3600
                if age_hours > max_age_hours:
                    old_sessions.append(session_id)
            
            for session_id in old_sessions:
                self.cleanup_session(session_id)
            
            if old_sessions:
                logger.info(f"清理了 {len(old_sessions)} 个旧会话")
                
        except Exception as e:
            logger.error(f"清理旧会话失败: {e}")
    
    def _calculate_total_progress(self, stage: ProgressStage, stage_progress: float) -> float:
        """
        计算总体进度百分比
        
        Args:
            stage: 当前阶段
            stage_progress: 阶段内进度 (0.0-1.0)
            
        Returns:
            float: 总体进度百分比 (0.0-100.0)
        """
        # 计算已完成阶段的权重
        completed_weight = 0
        for s in ProgressStage:
            if s == stage:
                break
            if s in self.stage_weights:
                completed_weight += self.stage_weights[s]
        
        # 计算当前阶段的权重
        current_stage_weight = self.stage_weights.get(stage, 0)
        current_progress = current_stage_weight * max(0, min(1, stage_progress))
        
        # 计算总进度
        total_progress = completed_weight + current_progress
        return total_progress
    
    def _stage_to_step(self, stage: ProgressStage) -> GenerationStep:
        """
        将进度阶段转换为生成步骤
        
        Args:
            stage: 进度阶段
            
        Returns:
            GenerationStep: 生成步骤
        """
        stage_mapping = {
            ProgressStage.INITIALIZING: GenerationStep.INITIALIZING,
            ProgressStage.GENERATING_OUTLINE: GenerationStep.GENERATING_OUTLINE,
            ProgressStage.CREATING_CONTENT: GenerationStep.CREATING_SLIDES,
            ProgressStage.APPLYING_TEMPLATE: GenerationStep.APPLYING_TEMPLATE,
            ProgressStage.GENERATING_PPT: GenerationStep.FINALIZING,
            ProgressStage.FINALIZING: GenerationStep.FINALIZING,
            ProgressStage.COMPLETED: GenerationStep.COMPLETED,
            ProgressStage.ERROR: GenerationStep.INITIALIZING
        }
        return stage_mapping.get(stage, GenerationStep.INITIALIZING)
    
    async def _send_progress_update(
        self,
        session_id: str,
        session: GenerationSession,
        message: Optional[str] = None,
        details: Optional[Dict[str, Any]] = None
    ):
        """
        发送进度更新消息
        
        Args:
            session_id: 会话ID
            session: 会话对象
            message: 进度消息
            details: 额外详情
        """
        try:
            # 构建进度更新消息
            progress_data = {
                "type": "progress_update",
                "session_id": session_id,
                "status": session.status.value,
                "progress": session.progress,
                "current_step": session.current_step.value,
                "message": message or f"进度: {session.progress}%",
                "updated_at": session.updated_at.isoformat()
            }
            
            # 添加额外详情
            if details:
                progress_data["details"] = details
            
            # 发送WebSocket消息
            await connection_manager.send_personal_message(session_id, progress_data)
            
        except Exception as e:
            logger.error(f"发送进度更新失败: {e}")


# 全局进度跟踪器实例
progress_tracker = ProgressTracker()